import subprocess
import pprint
import logging
import re
import xml.etree.ElementTree as etree

from config import *

logger = logging.getLogger('dmzwatch.'+__name__)


REPORT_FORMAT_PDF  = 'PDF'
REPORT_FORMAT_HTML = 'HTML'
REPORT_FORMAT_XML  = 'XML'


class TaskStatus(object):
    new = 'New'
    done = 'Done'
    requested = 'Requested'
    error = 'Internal Error'
    running = 'Running'
    
class OMPClient(object):
    def __init__ (self, host = MANAGER_ADDRESS, port = MANAGER_PORT, uname = MANAGER_USER, password = MANAGER_PASSWORD):
        self.host = host
        self.port = port
        self.uname = uname
        self.password = password
       
    def init_default_config(self, default_config, default_portlist):
        command = 'omp -h {} -p {} -u {} -w {} --get-configs'.format(self.host, self.port, self.uname, self.password)
        configs = subprocess.check_output(command, shell=True)
        for config in configs.decode('utf-8').split('\n')[:-1]:
            if config.split(None, 1)[1] == default_config:
                self.defaultconfig = OMPConfig(config.split(None, 1)[0],config.split(None, 1)[1])
                break

        command = 'omp -h {} -p {} -u {} -w {} -X \"<get_port_lists/>\"'.format(self.host, self.port, self.uname, self.password)
        result = subprocess.check_output(command, shell=True).decode('utf-8')
        root = etree.fromstring(result)
        root = etree.fromstring(result)
        for portlist in root.findall('port_list'):
            if portlist.find('name').text == default_portlist:
                self.defaultportlist = OMPPortList(portlist.attrib['id'], default_portlist)
                
    def init_default_reports(self):
        command = 'omp -h {} -p {} -u {} -w {} --get-report-formats'.format(self.host, self.port, self.uname, self.password)        
        formats = subprocess.check_output(command, shell=True)        
        for format in formats.decode('utf-8').split('\n')[:-1]:
            if format.split(None, 1)[1] == REPORT_FORMAT_PDF:
                self.PDFReport = OMPReport(format.split(None, 1)[0], format.split(None, 1)[1])
            elif format.split(None, 1)[1] == REPORT_FORMAT_HTML:
                self.HTMLReport = OMPReport(format.split(None, 1)[0], format.split(None, 1)[1])
            elif format.split(None, 1)[1] == REPORT_FORMAT_XML:
                self.XMLReport = OMPReport(format.split(None, 1)[0], format.split(None, 1)[1])     
        

    def get_targets(self):
        result = []
        command = 'omp -h {} -p {} -u {} -w {} --get-targets'.format(self.host, self.port, self.uname, self.password)
        targets = subprocess.check_output(command, shell=True)
        for target in targets.decode('utf-8').split('\n')[:-1]:
            result.append(OMPTarget(uid = target.split(None, 1)[0], name = target.split(None, 1)[1]))
    
        return result
        
    def get_tasks(self):    
        result = []
        command = 'omp -h {} -p {} -u {} -w {} --get-tasks'.format(self.host, self.port, self.uname, self.password)
        tasks = subprocess.check_output(command, shell=True)
        for task in tasks.decode('utf-8').split('\n')[:-1]:
            result.append(OMPTask(uid = task.split()[0], status = task.split()[1], name = task.split()[-1]))
        
        return result
        
    def create_target(self, address, name, portlist_uid):
        xml = '"<create_target><name>{}</name><hosts>{}</hosts><port_list id=\\"{}\\"></port_list></create_target>"'\
                .format(name, address, portlist_uid)
                
        command = 'omp -h {} -p {} -u {} -w {} -X {}'.format(self.host, self.port, self.uname, self.password, xml)
        result = subprocess.check_output(command, shell=True).decode('utf-8')
    
        
        root = etree.fromstring(result)
        status = root.attrib['status'] 
        
        if status == '201' or status == '202' or status == '203':
            return root.attrib['id']
        else:
            raise Exception('Cannot create target')
        
    def create_task(self, name, target_uid, config_uid):
        result = []
        command = 'omp -h {} -p {} -u {} -w {} --create-task --name={} --config={} --target={}'\
                    .format(self.host, self.port, self.uname, self.password, name, config_uid, target_uid)
                    
        result = subprocess.check_output(command, shell=True)
        
        return result.decode('utf-8').split()[0]
        
    def run_task(self, task_uid):
        result = []
        command = 'omp -h {} -p {} -u {} -w {} --start-task {}'.format(self.host, self.port, self.uname, self.password, task_uid)
        result = subprocess.check_output(command, shell=True)
        return result.decode('utf-8').split()[0]
        
    def get_status(self, task_uid):
        result = []
        command = 'omp -h {} -p {} -u {} -w {} --get-tasks {}'.format(self.host, self.port, self.uname, self.password, task_uid)
        task = subprocess.check_output(command, shell=True)
        return task.decode('utf-8').split('\n')[0].split()[1]
    
        
    def get_reports(self, report_uid):
        command = 'omp -h {} -p {} -u {} -w {} --get-report {} -f {}'.format(self.host, self.port, self.uname, self.password, report_uid, self.XMLReport.uid)
        xml = subprocess.check_output(command, shell=True)
        command = 'omp -h {} -p {} -u {} -w {} --get-report {} -f {}'.format(self.host, self.port, self.uname, self.password, report_uid, self.PDFReport.uid)
        pdf = subprocess.check_output(command, shell=True)
        command = 'omp -h {} -p {} -u {} -w {} --get-report {} -f {}'.format(self.host, self.port, self.uname, self.password, report_uid, self.HTMLReport.uid)
        html = subprocess.check_output(command, shell=True)
        return {'xml' : xml, 'pdf' : pdf, 'html' : html}
       
class OMPGeneric(object):
    def __init__ (self, uid, name):
        self.uid = uid
        self.name = name
    
    def __repr__(self):
        return 'UUID: {}, Name: {}'.format(self.uid, self.name)
        
class OMPTarget(OMPGeneric):
    def __init__ (self, uid, name):
        super().__init__(uid, name)
    
class OMPTask(OMPGeneric):
    def __init__ (self, uid, name, status = TaskStatus.new ):
        super().__init__(uid, name)
        self.status = status
        
    def __repr__(self):
        return super().__repr__() + ', Status: {}'.format(self.status)
        

class OMPConfig(OMPGeneric):
    def __init__ (self, uid, name):
        super().__init__(uid, name)  
        
class OMPReport(OMPGeneric):
    def __init__ (self, uid, name):
        super().__init__(uid, name)        
        
class OMPPortList(OMPGeneric):
    def __init__ (self, uid, name):
        super().__init__(uid, name)  
        
        

if __name__ == '__main__':
    loglevel = logging.DEBUG
    logging.basicConfig(format = u'%(filename)s[LINE:%(lineno)d]# %(levelname)-8s [%(asctime)s]  %(message)s', level = loglevel)

    a = OMPClient()
    a.init_default_config()

    
